#上下文管理模块
import json
import time
import os
from typing import List, Dict, Optional
from datetime import datetime

class ContextManager:
    def __init__(self, max_context_length: int = 5, storage_dir: str = "chat_sessions"):
        self.sessions = {}
        self.max_context_length = max_context_length
        self.storage_dir = storage_dir
        # 禁用会话持久化，每次进入都是新会话
        # self._ensure_storage_dir()
        # self._load_all_sessions()
    
    def get_context(self, session_id: str) -> List[Dict]:
        """获取会话上下文"""
        if session_id not in self.sessions:
            return []
        
        session = self.sessions[session_id]
        return session.get('context', [])
    
    def update_context(self, session_id: str, question: str, answer: str, classification: Dict):
        """更新会话上下文"""
        if session_id not in self.sessions:
            self.sessions[session_id] = {
                'created_at': datetime.now().isoformat(),
                'last_updated': datetime.now().isoformat(),
                'context': [],
                'question_count': 0,
                'question_types': []
            }
        
        session = self.sessions[session_id]
        
        # 添加新的对话轮次
        turn = {
            'timestamp': datetime.now().isoformat(),
            'question': question,
            'answer': answer,
            'classification': classification,
            'turn_id': len(session['context']) + 1
        }
        
        session['context'].append(turn)
        session['last_updated'] = datetime.now().isoformat()
        session['question_count'] += 1
        session['question_types'].append(classification['label'])
        
        # 限制上下文长度
        if len(session['context']) > self.max_context_length:
            session['context'] = session['context'][-self.max_context_length:]
        
        # 禁用会话保存，不持久化会话
        # self._save_session(session_id)
        
        # 更新会话统计
        self._update_session_stats(session_id)
    
    def get_session_summary(self, session_id: str) -> Dict:
        """获取会话摘要"""
        if session_id not in self.sessions:
            return {}
        
        session = self.sessions[session_id]
        context = session.get('context', [])
        
        # 统计问题类型
        type_counts = {}
        for turn in context:
            label = turn['classification']['label']
            type_counts[label] = type_counts.get(label, 0) + 1
        
        # 生成摘要
        summary = {
            'session_id': session_id,
            'created_at': session['created_at'],
            'last_updated': session['last_updated'],
            'total_turns': len(context),
            'question_types': type_counts,
            'current_topic': self._extract_current_topic(context),
            'learning_progress': self._assess_learning_progress(context)
        }
        
        return summary
    
    def _extract_current_topic(self, context: List[Dict]) -> str:
        """提取当前学习主题"""
        if not context:
            return "未开始"
        
        # 简单的主题提取逻辑
        recent_questions = [turn['question'] for turn in context[-3:]]
        question_text = " ".join(recent_questions).lower()
        
        # 关键词匹配
        topics = {
            '生命周期': ['生命周期', 'lifetime', '作用域'],
            '所有权': ['所有权', 'ownership', 'move'],
            '借用': ['借用', 'borrow', '引用'],
            '错误处理': ['错误', 'error', 'panic', 'result'],
            '迭代器': ['迭代器', 'iterator', 'for'],
            '模块': ['模块', 'module', 'use', 'mod']
        }
        
        for topic, keywords in topics.items():
            if any(keyword in question_text for keyword in keywords):
                return topic
        
        return "综合学习"
    
    def _assess_learning_progress(self, context: List[Dict]) -> Dict:
        """评估学习进度"""
        if not context:
            return {'level': 'beginner', 'confidence': 0.0}
        
        # 基于问题类型和复杂度评估
        type_scores = {
            'definition': 1,
            'usage': 2,
            'error_debug': 3,
            'comparison': 4,
            'faq': 1
        }
        
        total_score = 0
        for turn in context:
            label = turn['classification']['label']
            total_score += type_scores.get(label, 1)
        
        avg_score = total_score / len(context)
        
        if avg_score >= 3.5:
            level = 'advanced'
            confidence = min(0.9, avg_score / 5.0)
        elif avg_score >= 2.0:
            level = 'intermediate'
            confidence = min(0.8, avg_score / 3.0)
        else:
            level = 'beginner'
            confidence = min(0.7, avg_score / 2.0)
        
        return {
            'level': level,
            'confidence': confidence,
            'score': avg_score
        }
    
    def _update_session_stats(self, session_id: str):
        """更新会话统计信息"""
        session = self.sessions[session_id]
        
        # 计算会话时长
        created_at = datetime.fromisoformat(session['created_at'])
        last_updated = datetime.fromisoformat(session['last_updated'])
        duration = (last_updated - created_at).total_seconds()
        
        session['duration_seconds'] = duration
        session['avg_turn_interval'] = duration / session['question_count'] if session['question_count'] > 0 else 0
    
    def clear_context(self, session_id: str):
        """清除会话上下文"""
        if session_id in self.sessions:
            self.sessions[session_id]['context'] = []
            self.sessions[session_id]['last_updated'] = datetime.now().isoformat()
            # 保存到文件
            self._save_session(session_id)
    
    def get_all_sessions(self) -> Dict:
        """获取所有会话信息"""
        return {
            session_id: self.get_session_summary(session_id)
            for session_id in self.sessions.keys()
        }
    
    def save_sessions(self, file_path: str):
        """保存会话数据到文件"""
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(self.sessions, f, ensure_ascii=False, indent=2)
    
    def load_sessions(self, file_path: str):
        """从文件加载会话数据"""
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                self.sessions = json.load(f)
        except FileNotFoundError:
            print(f"会话文件 {file_path} 不存在，使用空会话")
            self.sessions = {}
    
    def _ensure_storage_dir(self):
        """确保存储目录存在"""
        if not os.path.exists(self.storage_dir):
            os.makedirs(self.storage_dir)
    
    def _save_session(self, session_id: str):
        """保存单个会话到文件"""
        if session_id in self.sessions:
            session_file = os.path.join(self.storage_dir, f"{session_id}.json")
            try:
                with open(session_file, 'w', encoding='utf-8') as f:
                    json.dump(self.sessions[session_id], f, ensure_ascii=False, indent=2)
            except Exception as e:
                print(f"保存会话 {session_id} 失败: {e}")
    
    def _load_all_sessions(self):
        """加载所有会话文件"""
        if not os.path.exists(self.storage_dir):
            return
        
        for filename in os.listdir(self.storage_dir):
            if filename.endswith('.json'):
                session_id = filename[:-5]  # 移除.json后缀
                session_file = os.path.join(self.storage_dir, filename)
                try:
                    with open(session_file, 'r', encoding='utf-8') as f:
                        self.sessions[session_id] = json.load(f)
                except Exception as e:
                    print(f"加载会话 {session_id} 失败: {e}")
    
    def delete_session(self, session_id: str):
        """删除会话文件和内存数据"""
        # 删除内存中的数据
        if session_id in self.sessions:
            del self.sessions[session_id]
        
        # 删除文件
        session_file = os.path.join(self.storage_dir, f"{session_id}.json")
        if os.path.exists(session_file):
            try:
                os.remove(session_file)
            except Exception as e:
                print(f"删除会话文件 {session_id} 失败: {e}")
    
    def get_storage_stats(self) -> Dict:
        """获取存储统计信息"""
        if not os.path.exists(self.storage_dir):
            return {"total_sessions": 0, "total_size": 0, "storage_dir": self.storage_dir}
        
        total_sessions = 0
        total_size = 0
        
        for filename in os.listdir(self.storage_dir):
            if filename.endswith('.json'):
                total_sessions += 1
                file_path = os.path.join(self.storage_dir, filename)
                total_size += os.path.getsize(file_path)
        
        return {
            "total_sessions": total_sessions,
            "total_size": total_size,
            "total_size_mb": round(total_size / (1024 * 1024), 2),
            "storage_dir": self.storage_dir
        } 